package org.teremail.mailbox.search;


import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.lucene.analysis.SimpleAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Hit;
import org.apache.lucene.search.Hits;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.WildcardQuery;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.spans.SpanOrQuery;
import org.apache.lucene.search.spans.SpanQuery;
import org.apache.lucene.search.spans.SpanTermQuery;
import org.apache.lucene.store.Directory;
import org.teremail.MailSystemException;
import org.teremail.mailbox.EntryDO;
import org.teremail.mailbox.FolderDO;
import org.teremail.mailbox.Logical;
import org.teremail.mailbox.MailboxDAO;
import org.teremail.message.Content;
import org.teremail.message.EntityVisitorAdaptor;
import org.teremail.message.Header;
import org.teremail.message.Message;
import org.teremail.message.SimpleEntity;
import org.teremail.message.SimpleMessage;
import org.teremail.store.Store;

import com.google.inject.Inject;

/**
 * Lucene based search implementation.
 * 
 * @author mike
 *
 */
public class LuceneMailboxSearch implements MailboxSearch {

    private final MailboxDAO mailboxDAO;
    private final Directory dir;
    private final static String ID_FIELD = "id";
    private final static String BODY_FIELD = "body";

    @Inject
    public LuceneMailboxSearch(@Logical MailboxDAO mailboxDAO, Directory dir) {
        this.mailboxDAO = mailboxDAO;
        this.dir = dir;
    }
    
    /**
     * Walk a search term structure and build a list of query objects to 
     * be used to query the index.
     * 
     * @param st
     * @return
     */
    private static List<Query> getTerms(SearchTerm st) {
        final List<Query> tqs = new ArrayList<Query>();
        
        SearchTermVisitor v = new SearchTermVisitor() {
            public void visitBody(String key) {
                tqs.add(new TermQuery(new Term(BODY_FIELD, key)));
            }

            public void visitHeader(String name, String value) {
                tqs.add(new WildcardQuery(new Term(name, "*" + value + "*")));
            }
        };
        
        st.accept(v);
        
        return tqs;
    }

    /**
     * Searches the specified folder using the supplied search term.  Will first
     * require a DB query to fetch a list of the message ids in the folder, then
     * will use those as part of the query to the index.
     */
    public Collection<EntryDO> search(FolderDO f, SearchTerm st) {
        List<EntryDO> entries = mailboxDAO.getEntries(f);
        Map<String,EntryDO> id2uidMap = new HashMap<String,EntryDO>(entries.size());
        
        SpanQuery[] sqs = new SpanQuery[entries.size()];
        int i = 0;
        for (EntryDO entry : entries) {
            sqs[i] = new SpanTermQuery(new Term(ID_FIELD, entry.getMessageId()));
            id2uidMap.put(entry.getMessageId(), entry);
            i++;
        }
        Query idQ = new SpanOrQuery(sqs);
        BooleanQuery q = new BooleanQuery();
        q.add(idQ, Occur.MUST);
        for (Query query : getTerms(st)) {
            q.add(query, Occur.MUST);            
        }
        
        Set<EntryDO> results = new HashSet<EntryDO>();
        IndexReader r = null;
        try {
            r = IndexReader.open(dir);
            IndexSearcher s = new IndexSearcher(r);
            Hits hits = s.search(q);
            
            for (Iterator<?> it = hits.iterator(); it.hasNext();) {
                Hit hit = (Hit) it.next();
                String id = hit.get(ID_FIELD);
                EntryDO edo = id2uidMap.get(id);
                if (edo != null) {
                    results.add(edo);
                }
            }
            
        } catch (CorruptIndexException e) {
            throw new MailSystemException(e);
        } catch (IOException e) {
            throw new MailSystemException(e);
        } finally {
            if (r != null) {
                try {
                    r.close();
                } catch (IOException e) {}
            }
        }

        return results;
    }

    /**
     * Add the message to the index.
     */
    public void index(Message message, Store store) {
        
        final List<Content> content = new ArrayList<Content>();
        
        
        message.accept(new EntityVisitorAdaptor() {
            public void visitSimpleEntity(SimpleEntity simpleEntity) {
                content.add(simpleEntity.getContent());
            }

            public void visitSimpleMessage(SimpleMessage simpleMessage) {
                content.add(simpleMessage.getContent());
            }
        });
        
        IndexWriter writer = null;
        try {
            writer = new IndexWriter(dir, new SimpleAnalyzer());
            Document d = new Document();
            
            d.add(new Field(ID_FIELD, message.getId(), Field.Store.YES, 
                    Field.Index.UN_TOKENIZED));
            
            for (Header header : message.getHeader()) {
                d.add(new Field(header.getName().toLowerCase(), 
                        header.getValue().toLowerCase(), Field.Store.NO, 
                        Field.Index.UN_TOKENIZED));
            }
            
            for (Content c : content) {
                InputStreamReader r = 
                    new InputStreamReader(store.getInputStream(c.getId()));
                d.add(new Field(BODY_FIELD, r));
            }
            writer.addDocument(d);
            
        } catch (Exception e) {
            throw new MailSystemException(e);
        } finally {
            if (writer != null) {
                try {
                    writer.close();
                } catch (IOException e) {}
            }
        }
        
    }

}
